use crate::xvector::{dot_product, XVector};
use anyhow::anyhow;
use std::{
    ops::{Add, AddAssign, Mul},
    sync::mpsc,
    thread,
};

const NUM_THREAD_NUM: usize = 4;

// [[1, 2], [1, 2], [1, 2]] => [1, 2, 1, 2, 1, 2]
pub struct Matrix<T> {
    pub data: Vec<T>,
    pub row: usize,
    pub col: usize,
}

pub struct XMsgInput<T> {
    index: usize,
    row: XVector<T>,
    col: XVector<T>,
}

impl<T> XMsgInput<T> {
    pub fn new(index: usize, row: XVector<T>, col: XVector<T>) -> Self {
        Self { index, row, col }
    }
}

pub struct XMsgOutput<T> {
    index: usize,
    value: T,
}

pub struct XMsg<T> {
    input: XMsgInput<T>,
    sender: oneshot::Sender<XMsgOutput<T>>,
}

impl<T> XMsg<T> {
    pub fn new(input: XMsgInput<T>, sender: oneshot::Sender<XMsgOutput<T>>) -> Self {
        Self { input, sender }
    }
}

pub fn multiply<T>(a: &Matrix<T>, b: &Matrix<T>) -> anyhow::Result<Matrix<T>>
where
    T: Add<Output = T> + Mul<Output = T> + AddAssign + Copy + Default + Send + 'static,
{
    if a.col != b.row {
        return Err(anyhow!("Matrix multiply error, a.col != b.row"));
    }

    let senders = (0..NUM_THREAD_NUM)
        .map(|_| {
            let (tx, rx) = mpsc::channel::<XMsg<T>>();
            thread::spawn(move || {
                for msg in rx {
                    let value = dot_product(msg.input.row, msg.input.col)?;
                    if let Err(e) = msg.sender.send(XMsgOutput {
                        index: msg.input.index,
                        value,
                    }) {
                        eprintln!("Send msg({}) error", e)
                    }
                }
                Ok::<_, anyhow::Error>(())
            });
            tx
        })
        .collect::<Vec<_>>();

    // let mut data = Vec::with_capacity(a.row * b.col);
    let matrix_len = a.row * b.col;
    let mut data = vec![T::default(); matrix_len];
    let mut receivers = Vec::with_capacity(matrix_len);
    for i in 0..a.row {
        for j in 0..b.col {
            let row = XVector::new(&a.data[i * a.col..(i + 1) * a.col]);
            let col_data = b.data[j..]
                .iter()
                .step_by(b.col)
                .copied()
                .collect::<Vec<_>>();
            let col = XVector::new(col_data);
            let index = i * b.col + j;
            let input = XMsgInput::new(index, row, col);
            let (tx, rx) = oneshot::channel();
            let msg = XMsg::new(input, tx);
            if let Err(e) = senders[index % NUM_THREAD_NUM].send(msg) {
                eprintln!("sender send msg error: {}", e);
            }
            receivers.push(rx);
        }
    }

    for rx in receivers {
        let ans = rx.recv()?;
        data[ans.index] = ans.value;
    }

    Ok(Matrix {
        data,
        row: a.row,
        col: b.col,
    })
}

impl<T> Matrix<T> {
    pub fn new(data: impl Into<Vec<T>>, row: usize, col: usize) -> Self {
        Self {
            data: data.into(),
            row,
            col,
        }
    }
}

impl<T> Mul for Matrix<T>
where
    T: Add<Output = T> + Mul<Output = T> + AddAssign + Copy + Default + Send + 'static,
{
    type Output = Self;
    fn mul(self, rhs: Self) -> Self::Output {
        multiply(&self, &rhs).expect("Matrix multiply error")
    }
}

impl<T> std::fmt::Display for Matrix<T>
where
    T: std::fmt::Display,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{{")?;
        for i in 0..self.row {
            for j in 0..self.col {
                write!(f, "{}", self.data[i * self.col + j])?;
                if j != self.col - 1 {
                    write!(f, " ")?;
                }
            }

            if i != self.row - 1 {
                write!(f, ", ")?;
            }
        }
        write!(f, "}}")?;
        Ok(())
    }
}

impl<T> std::fmt::Debug for Matrix<T>
where
    T: std::fmt::Display,
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "Matrix(row: {}, col: {}, data: {})",
            self.row, self.col, self
        )
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn matrix_multiply_test() -> anyhow::Result<()> {
        let a = Matrix::new([1, 2, 3, 4, 5, 6], 2, 3);
        let b = Matrix::new([1, 2, 3, 4, 5, 6], 3, 2);
        let c = a * b;
        assert_eq!(c.col, 2);
        assert_eq!(c.row, 2);
        assert_eq!(c.data, vec![22, 28, 49, 64]);
        assert_eq!(
            format!("{:?}", c),
            "Matrix(row: 2, col: 2, data: {22 28, 49 64})"
        );
        Ok(())
    }

    #[test]
    fn matrix_multiply_sample_test() -> anyhow::Result<()> {
        let a = Matrix::new([1, 2, 3, 4], 2, 2);
        let b = Matrix::new([1, 2, 3, 4], 2, 2);
        let c = a * b;
        assert_eq!(c.data, vec![7, 10, 15, 22]);
        assert_eq!(format!("{}", c), "{7 10, 15 22}");
        Ok(())
    }

    #[test]
    fn matrix_multiply_error_test() {
        let a = Matrix::new([1, 2, 3, 4, 5, 6], 2, 3);
        let b = Matrix::new([1, 2, 3, 4], 2, 2);
        let c = multiply(&a, &b);
        assert!(c.is_err())
    }

    #[test]
    #[should_panic]
    fn matrix_multiply_panic_test() {
        let a = Matrix::new([1, 2, 3, 4, 5, 6], 2, 3);
        let b = Matrix::new([1, 2, 3, 4], 2, 2);
        let _c = a * b;
    }
}
